import numpy as np
import matplotlib.pyplot as plt
from scipy.integrate import odeint

# 设置中文显示
plt.rcParams['font.sans-serif'] = ['SimHei']  # 用来正常显示中文标签
plt.rcParams['axes.unicode_minus'] = False  # 用来正常显示负号

# 固定参数
klist = [0.0022, 0.1369, 0.1754, 0.1616, 0.3022]
khlist = [0.2989, 0.6281, 0.8579, 0.7592]
# bound = np.array([[30, 50], [50, 90], [6, 10]])  # 温度, 湿度, 固含量的边界
bound = np.array([[30 + np.random.uniform(0, 2), 50 + np.random.uniform(-2, 0)], [50 + np.random.uniform(0, 2), 90 + np.random.uniform(-1, 0)], [6 + np.random.uniform(0, 1), 10 + np.random.uniform(-1, 0)]])  # 温度, 湿度, 固含量的边界



# 定义孔面积占比计算函数
def f(T, H, SC, klist, khlist):
    k_1, k_3, k_4, k_5, k_6 = klist
    k_H, S_C, S_S, eta_0 = khlist

    solid_content = SC / 100
    T_k = T + 273.15  # 转换为开尔文温度

    # 常量定义
    C_0 = 24e-3
    m_D_0 = 24e-3
    m_S = 6e-3
    m_C = (m_D_0 + m_S) * solid_content
    ro_D, ro_S, ro_C = 0.948e3, 1.261e3, 1.3e3
    V = m_D_0 / ro_D + m_S / ro_S + m_C / ro_C
    A, B, C = 6.09451, 2725.96, 28.209
    A_0 = -k_1 * (10 ** (A - B / (T_k + C)) * 133.322 * (1 - k_H * H / 100) / ro_D / V)
    n_cf = 6
    k_boltz = 1.380649e-23
    r = 0.3413e-9
    N_A = 6.02214076e23

    # 随时间变化的函数

    def m_D(t):
        return C_0 * np.exp(A_0 * t)

    def m_S_out(t):
        return max(0, m_S - m_D(t) * S_S)

    def m_C_out(t):
        return max(0, m_C - m_D(t) * S_C)

    def phi_C_out(t):
        return m_C_out(t) / ro_C / V

    def eta(temp, t):
        return eta_0 * (1 + 2.5 * phi_C_out(t))

    def D(temp, t):
        return k_boltz * temp / (n_cf * np.pi * r * eta(temp, t))

    def v(temp, t):
        return k_3 * (D(temp, t) ** 0.5)

    def n_density(t):
        return m_S_out(t) / (ro_S * 4 / 3 * np.pi * r ** 3) / V

    def Z(temp, t):
        return (2 ** 0.5) * n_density(t) * np.pi * (r ** 2) * v(temp, t)

    # 确定蒸发结束时间
    tf = 0
    for t in np.arange(0, 500, 0.01):
        if m_D(t) < 0.024 * 0.1:
            tf = t
            break

    # 解微分方程
    dy = lambda m_s, t: (k_4 * Z(T_k, t) + k_5 * v(T_k, t) * (m_S_out(t) - m_s) / V * m_s)
    t_points = np.arange(0.1, tf, 0.01)
    sol = odeint(dy, 0, t_points)

    return k_6 * sol[-1][0] if len(sol) > 0 else 0


# PSO算法实现
def PSO(objective_func, bounds, num_particles=30, max_iter=100):
    dim = len(bounds)
    particles = np.random.rand(num_particles, dim)
    lb, ub = bounds[:, 0], bounds[:, 1]
    particles = lb + particles * (ub - lb)  # 初始化位置

    velocities = np.zeros((num_particles, dim))
    personal_best_positions = np.copy(particles)
    personal_best_scores = np.array([objective_func(*p) for p in particles])

    global_best_idx = np.argmax(personal_best_scores)
    global_best_position = personal_best_positions[global_best_idx]
    global_best_score = personal_best_scores[global_best_idx]

    # 记录优化过程
    history = {
        'global_best_score': [],
        'global_best_position': [],
        'particles': [particles.copy()],
        'scores': [personal_best_scores.copy()],
        'temperature': [],  # 记录温度历史
        'humidity': [],  # 记录湿度历史
        'solid_content': [],  # 记录固含量历史
        'all_scores': []  # 记录所有粒子的分数
    }

    # 记录初始状态
    history['global_best_score'].append(global_best_score)
    history['global_best_position'].append(global_best_position)
    history['temperature'].append(global_best_position[0])
    history['humidity'].append(global_best_position[1])
    history['solid_content'].append(global_best_position[2])
    history['all_scores'].append(personal_best_scores.copy())

    # PSO参数
    w = 2  # 惯性权重
    c1, c2 = 0.2, 0.2  # 学习因子

    for iter in range(max_iter):
        r1, r2 = np.random.rand(2)
        for i in range(num_particles):
            # 更新速度
            velocities[i] = w * velocities[i] + \
                            c1 * r1 * (personal_best_positions[i] - particles[i]) + \
                            c2 * r2 * (global_best_position - particles[i])

            # 更新位置
            particles[i] += velocities[i]

            # 边界处理
            particles[i] = np.clip(particles[i], lb, ub)

            # 评估新位置
            current_score = objective_func(*particles[i])

            # 更新个体最优
            if current_score > personal_best_scores[i]:
                personal_best_scores[i] = current_score
                personal_best_positions[i] = particles[i].copy()

                # 更新全局最优
                if current_score > global_best_score:
                    global_best_score = current_score
                    global_best_position = particles[i].copy()

        # 记录当前迭代信息
        history['global_best_score'].append(global_best_score)
        history['global_best_position'].append(global_best_position)
        history['particles'].append(particles.copy())
        history['scores'].append(personal_best_scores.copy())
        history['temperature'].append(global_best_position[0])
        history['humidity'].append(global_best_position[1])
        history['solid_content'].append(global_best_position[2])
        history['all_scores'].append(personal_best_scores.copy())

        print(f'迭代 {iter + 1}/{max_iter}, 最佳孔面积占比: {global_best_score:.4f}%')

    return global_best_position, global_best_score, history


# 包装目标函数
def objective_func(T, H, SC):
    return f(T, H, SC, klist, khlist)


# 运行PSO优化
best_position, best_score, history = PSO(objective_func, bound, num_particles=50, max_iter=10)


# 可视化优化过程
def plot_PSO_history(history):
    plt.figure(figsize=(14, 10))
    plt.suptitle('多孔膜制备条件优化过程 - PSO算法', fontsize=16)
    plt.rc('font', size=16)
    plt.rc('font', family='SimHei')

    # 子图1：损失函数变化
    plt.subplot(2, 2, 1)
    plt.semilogy(history['global_best_score'], 'r-', label='全局最优解')
    plt.xlabel('迭代次数')
    plt.ylabel('孔面积占比 (%)')
    plt.title('孔面积占比优化过程')
    plt.legend()
    plt.grid(True, which="both", ls="--")

    # 子图2：温度衰减曲线
    plt.subplot(2, 2, 2)
    plt.plot(history['temperature'], 'g-')
    plt.plot(history['humidity'], 'b-')
    plt.plot(history['solid_content'], 'r-')
    plt.xlabel('迭代次数')
    plt.ylabel('温度、湿度、固含量')
    plt.title('制备条件曲线')
    plt.grid(True, ls="--")

    plt.tight_layout()
    plt.show()


# 绘制优化过程
plot_PSO_history(history)

# 输出最优解
print("\n最优制备条件:")
print(f"温度: {best_position[0]:.2f} °C")
print(f"湿度: {best_position[1]:.2f} %")
print(f"固含量: {best_position[2]:.2f} %")
print(f"最大孔面积占比: {best_score:.4f} %")

# 绘制最终参数分布
plt.figure(figsize=(10, 8))
final_particles = history['particles'][-1]
final_scores = history['scores'][-1]

# 创建3D图
ax = plt.subplot(111, projection='3d')

# 绘制所有粒子
scatter = ax.scatter(final_particles[:, 0], final_particles[:, 1], final_particles[:, 2],
                     c=final_scores, cmap='viridis', s=50, alpha=0.8)

# 绘制最优解
ax.scatter(best_position[0], best_position[1], best_position[2],
           c='red', s=200, marker='*', label='最优解')

ax.set_title('最终粒子分布')
ax.set_xlabel('温度 (°C)')
ax.set_ylabel('湿度 (%)')
ax.set_zlabel('固含量 (%)')

# 添加颜色条
cbar = plt.colorbar(scatter, ax=ax, pad=0.1)
cbar.set_label('孔面积占比 (%)')

plt.legend()
plt.tight_layout()
plt.show()